#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
列值统计工具
版本: 1.0.0
作者: 陈振玺
功能: 统计Excel/CSV/TSV文件中指定列的值出现次数
"""

import argparse
import pandas as pd
from pathlib import Path
import operator

class ColumnCounter:
    def __init__(self, input_file, column_name, output_file, cluster_columns=None, filter_column=None, filter_operator=None, filter_value=None):
        self.input_file = Path(input_file)
        self.column_name = column_name
        self.output_file = Path(output_file)
        self.cluster_columns = cluster_columns.split() if cluster_columns else None
        self.filter_column = filter_column
        self.filter_operator = filter_operator
        self.filter_value = filter_value
        self.data = None
        
        # 操作符映射
        self.operators = {
            '>': operator.gt,
            '<': operator.lt,
            '>=': operator.ge,
            '<=': operator.le,
            '==': operator.eq,
            '!=': operator.ne
        }
        
    def load_data(self):
        """加载数据文件"""
        try:
            file_extension = self.input_file.suffix.lower()
            
            if file_extension == '.xlsx':
                self.data = pd.read_excel(self.input_file)
            elif file_extension == '.csv':
                self.data = pd.read_csv(self.input_file)
            elif file_extension == '.tsv':
                self.data = pd.read_csv(self.input_file, sep='\t')
            else:
                raise ValueError(f"不支持的文件格式: {file_extension}")
                
            print(f"成功加载数据: {len(self.data)} 行, {len(self.data.columns)} 列")
            return True
            
        except Exception as e:
            print(f"加载数据失败: {e}")
            return False
    
    def validate_columns(self):
        """验证列名是否存在"""
        # 验证统计列
        if self.column_name not in self.data.columns:
            print(f"错误: 统计列 '{self.column_name}' 不存在")
            print(f"可用列名: {list(self.data.columns)}")
            return False
        
        # 验证聚类列
        if self.cluster_columns:
            for col in self.cluster_columns:
                if col not in self.data.columns:
                    print(f"错误: 聚类列 '{col}' 不存在")
                    print(f"可用列名: {list(self.data.columns)}")
                    return False
        
        # 验证过滤列
        if self.filter_column:
            if self.filter_column not in self.data.columns:
                print(f"错误: 过滤列 '{self.filter_column}' 不存在")
                print(f"可用列名: {list(self.data.columns)}")
                return False
            
            # 检查过滤列是否为数值型
            if not pd.api.types.is_numeric_dtype(self.data[self.filter_column]):
                print(f"错误: 过滤列 '{self.filter_column}' 不是数值型")
                return False
        
        return True
    
    def apply_filter(self):
        """应用数值过滤条件"""
        if not self.filter_column or not self.filter_operator or self.filter_value is None:
            return
        
        try:
            # 转换过滤值为数值
            filter_val = float(self.filter_value)
            
            # 获取操作符函数
            op_func = self.operators.get(self.filter_operator)
            if not op_func:
                print(f"错误: 不支持的操作符 '{self.filter_operator}'")
                print(f"支持的操作符: {list(self.operators.keys())}")
                return
            
            # 应用过滤条件
            original_count = len(self.data)
            self.data = self.data[op_func(self.data[self.filter_column], filter_val)]
            filtered_count = len(self.data)
            
            print(f"过滤条件: {self.filter_column} {self.filter_operator} {filter_val}")
            print(f"过滤结果: {original_count} -> {filtered_count} 行")
            
        except ValueError:
            print(f"错误: 过滤值 '{self.filter_value}' 不是有效数值")
        except Exception as e:
            print(f"过滤失败: {e}")
    
    def count_values(self):
        """统计指定列的值出现次数"""
        try:
            # 应用过滤条件
            self.apply_filter()
            
            if len(self.data) == 0:
                print("警告: 过滤后没有数据")
                return pd.DataFrame({'Item': [], 'Count': []})
            
            if self.cluster_columns:
                # 使用聚类列进行分组统计
                # 按聚类列和统计列进行分组计数
                group_cols = self.cluster_columns + [self.column_name]
                grouped = self.data.groupby(group_cols).size().reset_index(name='Count')
                
                # 按聚类列进行汇总，保留聚类列的值
                result_df = grouped.groupby(self.cluster_columns)['Count'].sum().reset_index()
                
                # 创建Item列：将聚类列的值用空格连接
                result_df['Item'] = result_df[self.cluster_columns].apply(
                    lambda x: ' '.join(x.astype(str)), axis=1
                )
                
                # 重新排列列的顺序：聚类列 + Item + Count
                cols = self.cluster_columns + ['Item', 'Count']
                result_df = result_df[cols]
                
                print(f"聚类统计完成: 使用聚类列 {self.cluster_columns}")
                print(f"发现 {len(result_df)} 个不同的聚类组合")
            else:
                # 普通统计值出现次数
                value_counts = self.data[self.column_name].value_counts()
                
                # 创建结果DataFrame
                result_df = pd.DataFrame({
                    'Item': value_counts.index,
                    'Count': value_counts.values
                })
                
                # 重置索引
                result_df = result_df.reset_index(drop=True)
                
                print(f"统计完成: 发现 {len(result_df)} 个不同的值")
            
            return result_df
            
        except Exception as e:
            print(f"统计失败: {e}")
            return None
    
    def save_result(self, result_df):
        """保存统计结果"""
        try:
            # 确保输出目录存在
            self.output_file.parent.mkdir(parents=True, exist_ok=True)
            
            # 根据输出文件扩展名选择保存格式
            output_extension = self.output_file.suffix.lower()
            
            if output_extension == '.xlsx':
                result_df.to_excel(self.output_file, index=False)
            elif output_extension == '.csv':
                result_df.to_csv(self.output_file, index=False)
            elif output_extension == '.tsv':
                result_df.to_csv(self.output_file, sep='\t', index=False)
            else:
                # 默认保存为CSV格式
                result_df.to_csv(self.output_file, index=False)
            
            print(f"结果已保存到: {self.output_file}")
            return True
            
        except Exception as e:
            print(f"保存失败: {e}")
            return False
    
    def process(self):
        """执行完整的统计流程"""
        print("开始列值统计...")
        
        # 1. 加载数据
        if not self.load_data():
            return False
        
        # 2. 验证列名
        if not self.validate_columns():
            return False
        
        # 3. 统计值出现次数
        result_df = self.count_values()
        if result_df is None:
            return False
        
        # 4. 保存结果
        if not self.save_result(result_df):
            return False
        
        print("\n✅ 列值统计完成!")
        print(f"📁 输入文件: {self.input_file}")
        print(f"📊 统计列名: {self.column_name}")
        print(f"📄 输出文件: {self.output_file}")
        print(f"📈 统计结果: {len(result_df)} 个不同值")
        
        return True

def main():
    parser = argparse.ArgumentParser(
        description='列值统计工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 基本统计
  column-counter input.csv status -o output.csv
  
  # 使用聚类列统计
  column-counter input.csv status -o output.csv --cluster "group category"
  
  # 添加过滤条件
  column-counter input.csv status -o output.csv --filter-column score --filter-op ">" --filter-value 80
  
  # 组合使用
  column-counter input.csv status -o output.csv --cluster "group category" --filter-column score --filter-op ">=" --filter-value 60

支持的过滤操作符: >, <, >=, <=, ==, !=
        """
    )
    parser.add_argument('input_file', help='输入文件路径 (支持Excel/CSV/TSV格式)')
    parser.add_argument('column_name', help='要统计的列名')
    parser.add_argument('-o', '--output', required=True, help='输出文件路径')
    parser.add_argument('--cluster', help='聚类列名，多个列名用空格分隔')
    parser.add_argument('--filter-column', help='用于过滤的数值型列名')
    parser.add_argument('--filter-op', choices=['>', '<', '>=', '<=', '==', '!='], help='过滤操作符')
    parser.add_argument('--filter-value', help='过滤值')
    
    args = parser.parse_args()
    
    # 验证过滤参数的完整性
    filter_params = [args.filter_column, args.filter_op, args.filter_value]
    if any(filter_params) and not all(filter_params):
        print("错误: 使用过滤功能时，必须同时指定 --filter-column, --filter-op, --filter-value")
        return 1
    
    counter = ColumnCounter(
        args.input_file, 
        args.column_name, 
        args.output,
        args.cluster,
        args.filter_column,
        args.filter_op,
        args.filter_value
    )
    success = counter.process()
    
    if success:
        print("\n🎉 列值统计工具执行成功!")
    else:
        print("\n❌ 列值统计工具执行失败!")
        exit(1)

if __name__ == '__main__':
    main()